import os
import json
import base64


data_labels_dirname = os.path.abspath(os.path.join(
    os.path.dirname(__file__), 'data_labels'
))
class LabelItem:
    def __init__(self, zh: str, en: str, count="0", alias: list[str]|None = None):
        self.zh = zh
        self.en = en
        self.count = count
        self.alias = alias


class LabelBase:
    """
    """
    
    CATEGORY = "destiny/Danbooru 标签"

    def __init__(self, label_map: dict):
        self.label_map = label_map

    RETURN_TYPES = ("STRING", )
    RETURN_NAMES = ("label",)

    FUNCTION = "execute"

    def execute(self, label: str, weight: float):
        item = self.label_map.get(label.split('-')[0], None)
        result = ""
        if item is not None:
            result = item.en
            if weight != 1:
                result = f"({result}:{round(weight, 2)})"
        return (result, )
    

def get_label_dict(file_name: str):
    filename = os.path.join(data_labels_dirname, file_name)
    if not os.path.exists(filename):
        return {}
    output = {}
    with open(filename, mode='r+', encoding='utf-8') as fpr:
        for line in fpr.readlines():
            line = line.strip()
            try:
                json_data = json.loads(base64.b64decode(line).decode('utf-8'))
            except:
                pass
            output[json_data['zh']] = LabelItem(
                zh=json_data['zh'], en=json_data['en'], count=str(json_data['count_number']), alias=json_data['alias_str']
            )
    return output


# ==== 以下是标签
# - 场景/室内
scene_indoor = get_label_dict('scene_indoor.txt')
class LabelSceneIndoor(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/场景"

    def __init__(self):
        super().__init__(scene_indoor)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in scene_indoor.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }


# - 人物/上身装饰
person_up_adorn = get_label_dict('person_up_adorn.txt')
class LabelPersonUpAdorn(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/人物"

    def __init__(self):
        super().__init__(person_up_adorn)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in person_up_adorn.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }

# - 人物/下身装饰
person_down_adorn = get_label_dict('person_down_adorn.txt')
class LabelPersonDownAdorn(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/人物"

    def __init__(self):
        super().__init__(person_down_adorn)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in person_down_adorn.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }

# - 人物/全身装饰
person_full_adorn = get_label_dict('person_full_adorn.txt')
class LabelPersonFullAdorn(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/人物"

    def __init__(self):
        super().__init__(person_full_adorn)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in person_full_adorn.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }

# - 人物/动作
person_action = get_label_dict('person_action.txt')
class LabelPersonAction(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/人物"

    def __init__(self):
        super().__init__(person_action)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in person_action.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }

# - 人物/品控
person_quality = get_label_dict('person_quality.txt')
class LabelPersonQuality(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/人物"

    def __init__(self):
        super().__init__(person_quality)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in person_quality.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }

# - 人物/头发
person_hair = get_label_dict('person_hair.txt')
class LabelPersonHair(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/人物"

    def __init__(self):
        super().__init__(person_hair)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in person_hair.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }

# - 人物/头部饰品
person_head_adorn = get_label_dict('person_head_adorn.txt')
class LabelPersonHeadAdorn(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/人物"

    def __init__(self):
        super().__init__(person_head_adorn)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in person_head_adorn.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }

# - 人物/姿势
person_pose = get_label_dict('person_pose.txt')
class LabelPersonPose(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/人物"

    def __init__(self):
        super().__init__(person_pose)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in person_pose.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }

# - 人物/类型
person_type = get_label_dict('person_type.txt')
class LabelPersonType(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/人物"

    def __init__(self):
        super().__init__(person_type)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in person_type.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }

# - 人物/耳朵
person_ear = get_label_dict('person_ear.txt')
class LabelPersonEar(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/人物"

    def __init__(self):
        super().__init__(person_ear)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in person_ear.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }

# - 人物/胸部
person_breasts = get_label_dict('person_breasts.txt')
class LabelPersonBreasts(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/人物"

    def __init__(self):
        super().__init__(person_breasts)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in person_breasts.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }

# - 人物/脖子
person_neck = get_label_dict('person_neck.txt')
class LabelPersonNeck(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/人物"

    def __init__(self):
        super().__init__(person_neck)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in person_neck.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }

# - 人物/衣装
person_clothing = get_label_dict('person_clothing.txt')
class LabelPersonClothing(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/人物"

    def __init__(self):
        super().__init__(person_clothing)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in person_clothing.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }

# - 人物/面部
person_face = get_label_dict('person_face.txt')
class LabelPersonFace(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/人物"

    def __init__(self):
        super().__init__(person_face)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in person_face.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }

# - 人物/鞋袜
person_shoes_socks = get_label_dict('person_shoes_socks.txt')
class LabelPersonShoesSocks(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/人物"

    def __init__(self):
        super().__init__(person_shoes_socks)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in person_shoes_socks.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }


# - 限制级/成年限定
restricted_adult = get_label_dict('restricted_adult.txt')
class LabelRestrictedAdult(LabelBase):
    """"""
    CATEGORY = "destiny/Danbooru 标签/限制级"

    def __init__(self):
        super().__init__(restricted_adult)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "label": ([f"{item.zh}-[{item.en}/{item.count}]" for item in restricted_adult.values()], {
                    "default": ""
                }),
                "weight": ("FLOAT", {"default": 1.0, "step": 0.01})
            }
        }